from multiprocessing import Manager,Pool,Lock
# from matplotlib import pyplot as plt
import numpy as np
import traceback
import random
import time
import math
import os
import scene
import shape
import B_Spline


# 路径优化器
class Path_Optimize():

    def __init__(self, border, dense, barriers, start = [0,0], num = 1):
        # Public Read only params
        self.num = num # 并行任务数
        self.loop = 100
        self.visibility_weight = 7
        self.curvature_weight = 10
        self.experence_weight = 8
        self.remember_rate = 0.6
        self.anti_repeat_rate = 0.0
        self.reward = 5
        self.type = 's'
        self.scene = scene.Scene(border,dense,barriers,self.type, start)
        self.scene.calc_dict()
        self.path_len = len(self.scene.joint)

        # Shared params
        self.experence_dict = {}
        self.optimal_path = {}

        # Multiprocessing Tool
        self.lock = None

    # ----------------主进程子任务-----------------------------------------------
    def init_experence_dict(self):
        init_value = 1
        for i in range(self.path_len):
            for j in range(i+1,self.path_len):
                self.experence_dict['{}-{}'.format(i,j)] = init_value

    def init_optimal_path(self):
        self.optimal_path['score'] = 0
        self.optimal_path['path'] = []

    # ----------------子进程子任务-----------------------------------------------
    # 此时修改 Shared params 需加锁
    def experence_decay(self):
        try:
            #sum = 0 #--
            self.lock.acquire()
            for key in self.experence_dict.keys():
                self.experence_dict[key] *= self.remember_rate
                #sum += self.experence_dict[key] #--
            self.lock.release()
            # print('[{}]sum:{}'.format(os.getpid(),sum)) #--
        except:
            traceback.print_exc()

    def set_experence_dict(self,score,path):
        try:
            self.experence_decay()
            fix_coef = 1/50
            self.lock.acquire()
            for i in range(len(path)):
                cur = path[i]
                pre = path[i - 1]
                if (pre < cur):
                    self.experence_dict['{}-{}'.format(pre,cur)] += score*fix_coef
                elif (pre > cur):
                    self.experence_dict['{}-{}'.format(cur,pre)] += score*fix_coef
            self.lock.release()
        except:
            traceback.print_exc()   

    def set_optimal_path(self,score,path):
        try:
            best_score = self.optimal_path['score']
            if (score > best_score):
                #print('Best:{:.3f}'.format(score))
                self.lock.acquire()
                self.optimal_path['score'] = score
                self.optimal_path['path'] = path
                self.lock.release()
                if (best_score == 0): #首次不给奖励
                    return 0
                else: # 最优路径加奖励
                    return 1
            else:
                return 0
        except:
            traceback.print_exc() 
            return 0    

    def get_experence_list(self,current_point):
        try:
            experence_list = []
            for i in range(self.path_len):
                if (i == current_point):
                    experence_list.append(0.0)
                elif (i < current_point):
                    experence_list.append(self.experence_dict['{}-{}'.format(i,current_point)])
                else:
                    experence_list.append(self.experence_dict['{}-{}'.format(current_point,i)])
            return np.array(experence_list)
        except:
            traceback.print_exc()
            return np.array([0])

    def set_norm_prob_list(self,vis_list,curv_list,anti_repeat_list,experence_list):
        try:

            head_list = ((np.exp(vis_list*self.visibility_weight)-1)
                        *(np.exp(curv_list*self.curvature_weight)-1)* anti_repeat_list)
            prob_list = np.array(head_list * np.exp(experence_list*self.experence_weight))

            sum = np.sum(prob_list) 

            if (sum != 0):
                prob_list /= sum
                return prob_list
            else:
                return np.array([0])
        except:
            traceback.print_exc()
            return np.array([0])

    def choose_path(self,prob_list):
        try:
            arrow = random.uniform(0,1) # 概率箭
            
            target_disk = 0   # 靶盘
            for i in range(len(prob_list)): 
                target_disk += prob_list[i]
                if arrow <= target_disk:
                    # if i == len(prob_list)-1:
                    #     print('arrow:{:.3f},sum:{:.1f},tar:{}'.format(arrow,np.sum(prob_list),i))
                    return i
            # if i == len(prob_list)-1:
            #     print('end||arrow:{:.3f},sum:{:.1f},tar:{}'.format(arrow,np.sum(prob_list),i))
            #     print(prob_list)
            return i
        except:
            traceback.print_exc()

    def end_path(self,anti_repeat_list):
        try:
            anti_repeat_list = np.array(anti_repeat_list)
            max_anti_repeat = np.max(anti_repeat_list) ### stack overflow!!!
            # print('###',anti_repeat_list,max_anti_repeat)
            if (max_anti_repeat < 1):
                return True
            else:
                return False
        except:
            traceback.print_exc()
            return False


    # ----------------子进程主任务--------------------------------------------------
    def task(self,force_complete):
        try:
            start_time = time.time()
            # Private params
            current_point = 0 #random.randrange(0, self.path_len)# 优化路径若成环，则起点随机初始化
            post_point = current_point
            curv_list = np.ones(self.path_len)
            best_score = 0
            skip = 0

            for i in range(self.loop):
                score = 0
                current_point = 0
                post_point = 0
                path = [current_point]
                anti_repeat_list = np.ones(self.path_len)
                anti_repeat_list[current_point] *= self.anti_repeat_rate
                len_p = 0

                while(not self.end_path(anti_repeat_list)): ### stack overflow!!!
                    # Front update
                    vis_list = self.scene.sql_vis(current_point)
                    curv_list = self.scene.sql_curv(post_point,current_point)
                    experence_list = self.get_experence_list(current_point)
                    probility_list = self.set_norm_prob_list(vis_list,curv_list,anti_repeat_list,experence_list)
                    len_p = len(probility_list)
                    max_p = max(probility_list)
                    if (len_p < self.path_len or max_p == 0):
                        print('No choice, Quit-{}.'.format(max_p))
                        break

                    # Step forward
                    post_point = current_point
                    current_point = self.choose_path(probility_list)
                    anti_repeat_list[current_point] *= self.anti_repeat_rate
                    path.append(current_point)

                    # Back update
                    #step_score = math.pow(vis_list[current_point],self.visibility_weight)*math.pow(curv_list[current_point],self.curvature_weight)
                    step_score = vis_list[current_point]*self.visibility_weight + curv_list[current_point]*self.curvature_weight
                    score += step_score
                    #print('-->cur:{},s_score:{:.3f}'.format(current_point,step_score))

                if (len_p == force_complete):
                    skip += 1
                    continue
                score = score/(self.path_len + abs(len(path)-self.path_len))
                if (best_score < score):
                    best_score = score
                if(self.set_optimal_path(score,path) == 1):
                    score += self.reward
                self.set_experence_dict(score,path)

            if (best_score == 0 and force_complete == 1):
                self.task(0) 
            end_time = time.time()
            print('[{}]-Score:{:.3f},Skip:{},Runtime:{:.3f}'.format(os.getpid(),best_score,skip,end_time-start_time))

        except:
            traceback.print_exc()

    # ---------------子进程回调任务------------------------------------------------
    def err(self,err):
        print(err)

    # ----------------主进程主任务------------------------------------------------
    def run(self):
        try:
            optimal_path = []
            optimal_path_id = []
            if (self.num > 1):
                with Manager() as m:
                    print('---start---')
                    self.lock = m.Lock()
                    self.experence_dict = m.dict()
                    self.optimal_path = m.dict()
                    self.init_experence_dict()
                    self.init_optimal_path()
                    pool = Pool()
                    for i in range(self.num):
                        pool.apply_async(func=self.task,args=(1,),error_callback=self.err)
                    pool.close()
                    pool.join()
                    print('==> Final score:{:.3f}'.format(self.optimal_path['score']))
                    optimal_path_id = self.optimal_path['path']
            else:
                    self.lock = Lock()
                    self.experence_dict = {}
                    self.optimal_path = {}
                    self.init_experence_dict()
                    self.init_optimal_path()
                    self.task(1)
                    print('==> Final score:{:.3f}'.format(self.optimal_path['score']))
                    optimal_path_id = self.optimal_path['path']

            for id in optimal_path_id:
                optimal_path.append(self.scene.joint[id])
            print('==> Count loss:{}'.format(len(optimal_path_id) - self.path_len))

            return np.array(optimal_path)
        except:
            traceback.print_exc()

if not __name__ == '__main__':

    #for i in range(50):

        # start_time = time.time()

        # border = [-25,-25,500,325]
        # rpx = random.randrange(20, 480)
        # rpy = random.randrange(20, 300)
        # rpt = random.randrange(0, 180) - 89
        # pillar1 = shape.Pillar([100,140],45,100)
        # # pillar2 = shape.Pillar([390,110],-10,100)
        # # pillar3 = shape.Pillar([330,170],90,100)
        # pillar_random = shape.Pillar([rpx,rpy],rpt,100)
        # circle1 = shape.Circle([360,160],50)

        # mid_time = time.time()
        # op = Path_Optimize(border,[pillar_random],1)
        # optimal_path = op.run()
        # print('Optimize time:{:.3f}'.format(time.time()-mid_time))

        # px = op.scene.joint[:,:1].flatten()
        # py = op.scene.joint[:,1:].flatten()

        # # ox = optimal_path[:,:1].flatten()
        # # oy = optimal_path[:,1:].flatten()

        # mid_time = time.time()
        # bs = B_Spline.BSpline(optimal_path)
        # curve = bs.curve(500)
        # print('Curve time:{:.3f}'.format(time.time()-mid_time))


        # cx = curve[:,:1].flatten()
        # cy = curve[:,1:].flatten()

        # mid_time = time.time()
        # plt.title('Path generated in {:.3f}s'.format(mid_time-start_time))
        # for sp in op.scene.barriers:
        #     sp.draw()
        # plt.scatter(px,py,s=5,c='r')
        # #plt.plot(ox,oy)
        # plt.scatter(cx,cy,s=1,c='g')
        # plt.savefig('./pytool/path_result/path_{}.png'.format(time.gmtime().tm_sec), dpi=300)
        # print('Plot time:{:.3f}'.format(time.time()-mid_time))

        # end_time = time.time()
        # print('Running time: {:.3f}'.format(end_time - start_time))
        # #plt.show()
        # plt.clf()
    pass